# ascii_mode.py —— 支持 ASCII 字符串按 2 字节分组的加解密（补零对齐）
import saes

def bytes_to_blocks(b: bytes):
    if len(b) % 2:  # 奇数长度补 0
        b += b"\x00"
    for i in range(0, len(b), 2):
        yield (b[i] << 8) | b[i+1]

def blocks_to_bytes(blocks):
    out = bytearray()
    for w in blocks:
        out += bytes([(w >> 8) & 0xFF, w & 0xFF])
    return bytes(out)

def encrypt_ascii(s: str, key16: int) -> bytes:
    data = s.encode("ascii", errors="strict")
    c_blocks = [saes.encrypt_block16(w, key16) for w in bytes_to_blocks(data)]
    return blocks_to_bytes(c_blocks)

def decrypt_ascii(c: bytes, key16: int) -> str:
    assert len(c) % 2 == 0
    p_blocks = [saes.decrypt_block16((c[i]<<8)|c[i+1], key16) for i in range(0, len(c), 2)]
    p = blocks_to_bytes(p_blocks)
    return p.rstrip(b"\x00").decode("ascii", errors="strict")  # 去掉可能的补零

if __name__ == "__main__":
    K = 0x1F2E
    msg = "Hello S-AES!"
    ct = encrypt_ascii(msg, K)
    print("Cipher bytes:", ct.hex())
    pt = decrypt_ascii(ct, K)
    print("Plain:", pt)
